import asyncio
import os
import shutil
import tempfile
import resource
import signal
import psutil
import subprocess
import json
from datetime import datetime, timedelta
from typing import Dict, Any, List, Optional
from pathlib import Path
import logging

from .base import BaseExecutor
from ..entities import (
    CodeExecutionRequest, 
    ExecutionResult, 
    TaskStatus, 
    CodeType,
    ExecutorConfig
)

logger = logging.getLogger(__name__)

class SecureExecExecutor(BaseExecutor):
    """安全加固的代码执行器"""
    
    def __init__(self, config: ExecutorConfig):
        super().__init__(config)
        self.temp_dirs: Dict[str, str] = {}
        
        # 危险模块黑名单
        self.python_blocked_modules = {
            # 系统操作
            'os', 'sys', 'subprocess', 'shutil', 'pathlib',
            # 网络相关
            'socket', 'urllib', 'urllib2', 'urllib3', 'requests', 'http', 'httplib',
            'ftplib', 'smtplib', 'telnetlib', 'xmlrpc', 'webbrowser',
            # 文件系统
            'glob', 'tempfile', 'shutil', 'zipfile', 'tarfile',
            # 进程和线程
            'threading', 'multiprocessing', 'concurrent', 'asyncio',
            # 系统信息
            'platform', 'getpass', 'pwd', 'grp', 'spwd',
            # 危险内置函数
            'exec', 'eval', 'compile', '__import__', 'open', 'file',
            # 其他危险模块
            'ctypes', 'marshal', 'pickle', 'shelve', 'dbm',
            'sqlite3', 'codecs', 'encodings'
        }
        
        # JavaScript危险对象/函数黑名单
        self.js_blocked_features = [
            'require', 'process', 'Buffer', 'global', '__dirname', '__filename',
            'module', 'exports', 'eval', 'Function', 'setTimeout', 'setInterval',
            'clearTimeout', 'clearInterval'
        ]
        
        # 代码类型到执行命令的映射
        self.code_runners = {
            CodeType.PYTHON: self._get_python_runner,
            CodeType.JAVASCRIPT: self._get_nodejs_runner,
            CodeType.BASH: self._get_bash_runner,
        }
        
        # 文件扩展名映射
        self.file_extensions = {
            CodeType.PYTHON: ".py",
            CodeType.JAVASCRIPT: ".js", 
            CodeType.BASH: ".sh",
        }
    
    def _create_restricted_python_wrapper(self, original_code: str, code_file: str) -> str:
        """创建受限的Python代码包装器 - main函数调用模式"""
        wrapper_code = f'''
import sys
import builtins
import json

# 保存原始的__import__函数
original_import = builtins.__import__

# 重写 __import__ 函数以阻止危险模块
def safe_import(name, globals=None, locals=None, fromlist=(), level=0):
    blocked_modules = {self.python_blocked_modules}
    
    # 检查直接导入
    if name in blocked_modules:
        raise ImportError(f"模块 '{{name}}' 被禁用")
    
    # 检查子模块导入
    for blocked in blocked_modules:
        if name.startswith(blocked + '.'):
            raise ImportError(f"模块 '{{name}}' 被禁用")
    
    # 检查从模块导入
    if fromlist:
        for item in fromlist:
            if f"{{name}}.{{item}}" in blocked_modules:
                raise ImportError(f"模块 '{{name}}.{{item}}' 被禁用")
    
    return original_import(name, globals, locals, fromlist, level)

# 替换__import__函数
builtins.__import__ = safe_import

# 禁用危险的内置函数
def disabled_exec(*args, **kwargs):
    raise RuntimeError("exec() 被禁用")

def disabled_eval(*args, **kwargs):
    raise RuntimeError("eval() 被禁用")

def disabled_compile(*args, **kwargs):
    raise RuntimeError("compile() 被禁用")

def disabled_open(*args, **kwargs):
    raise RuntimeError("open() 被禁用，请使用print()输出")

# 替换危险函数
builtins.exec = disabled_exec
builtins.eval = disabled_eval
builtins.compile = disabled_compile
builtins.open = disabled_open

# 删除其他危险的内置变量（但保留__import__）
dangerous_builtins = ['__loader__', '__spec__']
for name in dangerous_builtins:
    if hasattr(builtins, name):
        try:
            delattr(builtins, name)
        except:
            pass

# 用户代码开始
{original_code}

# main函数调用和结果处理
if __name__ == "__main__":
    try:
        # 检查main函数是否存在
        if 'main' not in globals():
            print(json.dumps({{"error": "代码中未找到main函数"}}, ensure_ascii=False))
            sys.exit(1)
        
        # 获取输入参数
        input_args = {{}}
        if len(sys.argv) > 1:
            try:
                input_args = json.loads(sys.argv[1])
            except:
                # 兼容旧格式，如果不是JSON则作为字符串处理
                input_args = {{"input": sys.argv[1]}}
        
        # 调用main函数，传递解析后的参数
        result = main(**input_args)
        
        # 输出结果
        if result is not None:
            print(json.dumps(result, ensure_ascii=False, default=str))
        else:
            print(json.dumps({{}}, ensure_ascii=False))
            
    except Exception as e:
        # 输出错误信息
        error_info = {{
            "error": str(e),
            "error_type": type(e).__name__
        }}
        print(json.dumps(error_info, ensure_ascii=False))
        sys.exit(1)
'''
        return wrapper_code
    
    def _create_restricted_js_wrapper(self, original_code: str) -> str:
        """创建受限的JavaScript代码包装器 - main函数调用模式"""
        wrapper_code = f'''
// 安全限制
const originalRequire = typeof require !== 'undefined' ? require : undefined;
const blockedModules = ['fs', 'path', 'os', 'child_process', 'cluster', 'net', 'http', 'https', 'url', 'querystring', 'crypto'];

if (originalRequire) {{
    require = function(module) {{
        if (blockedModules.includes(module)) {{
            throw new Error(`模块 '${{module}}' 被禁用`);
        }}
        return originalRequire(module);
    }};
}}

// 禁用危险的全局对象
if (typeof global !== 'undefined') {{
    delete global.process;
    delete global.Buffer;
    delete global.__dirname;
    delete global.__filename;
}}

// 禁用危险函数
if (typeof eval !== 'undefined') {{
    eval = function() {{ throw new Error('eval() 被禁用'); }};
}}
if (typeof Function !== 'undefined') {{
    const OriginalFunction = Function;
    Function = function() {{ throw new Error('Function构造器被禁用'); }};
}}

// 用户代码
{original_code}

// main函数调用和结果处理
(function() {{
    try {{
        // 检查main函数是否存在
        if (typeof main !== 'function') {{
            console.log(JSON.stringify({{error: "代码中未找到main函数"}}));
            process.exit(1);
        }}
        
        // 获取输入参数
        let inputArgs = {{}};
        if (process.argv[2]) {{
            try {{
                inputArgs = JSON.parse(process.argv[2]);
            }} catch (e) {{
                // 兼容旧格式，如果不是JSON则作为字符串处理
                inputArgs = {{input: process.argv[2]}};
            }}
        }}
        
        // 调用main函数，传递解析后的参数
        const result = main(inputArgs);
        
        // 输出结果
        if (result !== undefined && result !== null) {{
            console.log(JSON.stringify(result));
        }} else {{
            console.log(JSON.stringify({{}}));
        }}
        
    }} catch (e) {{
        // 输出错误信息
        const errorInfo = {{
            error: e.message || String(e),
            error_type: e.constructor.name
        }};
        console.log(JSON.stringify(errorInfo));
        process.exit(1);
    }}
}})();
'''
        return wrapper_code
    
    def _create_restricted_bash_wrapper(self, original_code: str) -> str:
        """创建受限的Bash代码包装器 - main函数调用模式"""
        wrapper_code = f'''#!/bin/bash

# 设置严格模式
set -euo pipefail

# 限制环境变量
unset SHELL BASH BASH_ENV

# 用户代码
{original_code}

# main函数调用和结果处理
if ! declare -f main > /dev/null 2>&1; then
    echo '{{"error": "代码中未找到main函数"}}' 
    exit 1
fi

# 获取输入参数
input_json="${{1:-}}"

# 解析JSON参数并设置为环境变量（bash doesn't support objects natively）
if [ -n "$input_json" ] && [ "$input_json" != "{{}}" ]; then
    # 尝试解析JSON并导出变量
    if command -v jq > /dev/null 2>&1; then
        # 将JSON的每个key-value对设置为环境变量
        eval $(echo "$input_json" | jq -r 'to_entries[] | "export ARG_" + (.key | ascii_upcase) + "=" + (.value | tostring | @sh)')
    fi
fi

# 调用main函数并捕获结果
if result=$(main 2>&1); then
    if [ -n "$result" ]; then
        # 尝试格式化为JSON，如果不是JSON则包装
        if echo "$result" | jq . > /dev/null 2>&1; then
            echo "$result"
        else
            echo "{{\\"result\\": \\"$result\\"}}" | jq .
        fi
    else
        echo '{{}}'
    fi
else
    exit_code=$?
    echo '{{"error": "main函数执行失败", "exit_code":' "$exit_code" '}}'
    exit 1
fi
'''
        return wrapper_code
    
    def _check_main_function_exists(self, code: str, code_type: CodeType) -> bool:
        """检查代码中是否存在main函数"""
        if code_type == CodeType.PYTHON:
            # 简单检查Python中是否有main函数定义
            return 'def main(' in code
        elif code_type == CodeType.JAVASCRIPT:
            # 简单检查JavaScript中是否有main函数定义
            return 'function main(' in code or 'const main =' in code or 'let main =' in code or 'var main =' in code
        elif code_type == CodeType.BASH:
            # 简单检查Bash中是否有main函数定义
            return 'main()' in code or 'function main' in code
        return False
    
    def _get_python_runner(self, code_file: str, timeout_seconds: int = 30, input_arg: str = "") -> List[str]:
        """获取Python执行命令"""
        cmd = [
            "timeout", f"{timeout_seconds}s",
            "python3", "-u", "-S",  # -S: 不导入site模块
            code_file
        ]
        if input_arg:
            cmd.append(input_arg)
        return cmd
    
    def _get_nodejs_runner(self, code_file: str, timeout_seconds: int = 30, input_arg: str = "") -> List[str]:
        """获取Node.js执行命令（受限）"""
        cmd = [
            "timeout", f"{timeout_seconds}s",
            "node", "--max-old-space-size=128",
            code_file
        ]
        if input_arg:
            cmd.append(input_arg)
        return cmd
    
    def _get_bash_runner(self, code_file: str, timeout_seconds: int = 30, input_arg: str = "") -> List[str]:
        """获取Bash执行命令（受限）"""
        cmd = [
            "timeout", f"{timeout_seconds}s",
            "bash", code_file
        ]
        if input_arg:
            cmd.append(input_arg)
        return cmd
    
    async def prepare_environment(self, request: CodeExecutionRequest) -> Dict[str, Any]:
        """准备代码执行环境"""
        # 创建隔离的临时目录
        temp_dir = tempfile.mkdtemp(prefix="secure_sandbox_")
        self.temp_dirs[str(id(request))] = temp_dir
        
        # 创建受限的工作目录结构
        work_dir = os.path.join(temp_dir, "work")
        os.makedirs(work_dir, mode=0o700)
        
        # 获取文件扩展名
        extension = self.file_extensions.get(request.code_type, ".txt")
        code_file = os.path.join(work_dir, f"code{extension}")
        
        # 处理代码内容（添加安全包装）
        processed_code = self._process_code_for_security(request.code, request.code_type)
        
        # 写入代码文件
        with open(code_file, 'w', encoding='utf-8') as f:
            f.write(processed_code)
        
        # 设置严格的文件权限
        os.chmod(code_file, 0o600)  # 只有所有者可读写
        os.chmod(work_dir, 0o700)   # 只有所有者可访问目录
        
        logger.debug(f"安全代码已写入: {code_file}")
        
        return {
            "temp_dir": temp_dir,
            "work_dir": work_dir,
            "code_file": code_file,
            "original_code": request.code,
            "processed_code": processed_code
        }
    
    def _process_code_for_security(self, code: str, code_type: CodeType) -> str:
        """为安全性处理代码"""
        # 首先检查main函数是否存在
        if not self._check_main_function_exists(code, code_type):
            raise ValueError(f"代码中未找到main函数，请确保代码包含main函数定义")
        
        if code_type == CodeType.PYTHON:
            return self._create_restricted_python_wrapper(code, "")
        elif code_type == CodeType.JAVASCRIPT:
            return self._create_restricted_js_wrapper(code)
        elif code_type == CodeType.BASH:
            return self._create_restricted_bash_wrapper(code)
        else:
            return code
    
    def _create_restricted_bash_wrapper(self, original_code: str) -> str:
        """创建受限的Bash代码包装器 - 使用强化安全机制"""
        # 使用强化Bash安全执行器创建安全包装脚本
        return self.bash_executor.create_secure_bash_wrapper(original_code)
    
    async def apply_security_constraints(self, request: CodeExecutionRequest, env_info: Dict[str, Any]) -> Dict[str, Any]:
        """应用严格的安全约束"""
        
        # 严格的资源限制
        memory_limit = min(request.memory_limit_mb * 1024 * 1024, 64 * 1024 * 1024)  # 最大64MB
        cpu_time_limit = min(request.timeout_seconds, 10)  # 最大10秒
        file_size_limit = 1024 * 1024  # 1MB文件大小限制
        process_limit = 1  # 只允许一个进程
        
        security_config = {
            "memory_limit": memory_limit,
            "cpu_time_limit": cpu_time_limit,
            "file_size_limit": file_size_limit,
            "process_limit": process_limit,
            "working_dir": env_info["work_dir"],
            "network_disabled": True,
            "filesystem_readonly": True,
            "no_new_privileges": True
        }
        
        return security_config
    
    def _create_secure_preexec_fn(self, security_config: Dict[str, Any], code_type: Optional[CodeType] = None):
        """创建安全的进程启动前配置函数"""
        def preexec():
            try:
                # 创建新的进程组
                os.setpgrp()
                
                # 根据代码类型调整限制
                if code_type == CodeType.JAVASCRIPT:
                    # JavaScript需要更宽松的限制以支持Node.js运行时
                    memory_limit = max(security_config["memory_limit"], 512 * 1024 * 1024)  # 至少512MB
                    file_desc_limit = 100  # 更多文件描述符
                    process_limit = 15     # 允许更多进程
                else:
                    memory_limit = security_config["memory_limit"]
                    file_desc_limit = 10
                    process_limit = security_config["process_limit"]
                
                # 设置资源限制
                # 对JavaScript暂时不限制虚拟内存，因为Node.js启动需要大量虚拟内存
                if code_type != CodeType.JAVASCRIPT:
                    resource.setrlimit(resource.RLIMIT_AS, (memory_limit, memory_limit))
                
                resource.setrlimit(resource.RLIMIT_CPU, (
                    security_config["cpu_time_limit"], 
                    security_config["cpu_time_limit"]
                ))
                resource.setrlimit(resource.RLIMIT_FSIZE, (
                    security_config["file_size_limit"], 
                    security_config["file_size_limit"]
                ))
                resource.setrlimit(resource.RLIMIT_NPROC, (process_limit, process_limit))
                
                # 设置核心转储大小为0
                resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
                
                # 设置文件描述符限制
                try:
                    resource.setrlimit(resource.RLIMIT_NOFILE, (file_desc_limit, file_desc_limit))
                except (OSError, ValueError):
                    pass
                
            except Exception as e:
                logger.error(f"设置安全约束失败: {e}")
        
        return preexec
    
    def _create_secure_environment(self, request: CodeExecutionRequest, work_dir: str) -> Dict[str, str]:
        """创建安全的环境变量"""
        # 清理环境变量，只保留必要的
        secure_env = {
            'PATH': '/usr/local/bin:/usr/bin:/bin',  # 包含更多路径以支持Node.js
            'HOME': work_dir,
            'TMPDIR': work_dir,
            'TEMP': work_dir,
            'TMP': work_dir,
            'PWD': work_dir,
            'LANG': 'C.UTF-8',
            'LC_ALL': 'C.UTF-8',
            # 禁用Python字节码缓存
            'PYTHONDONTWRITEBYTECODE': '1',
            'PYTHONUSERBASE': '/dev/null',
            'PYTHONPATH': '',
        }
        
        # JavaScript特定的环境变量
        if request.code_type == CodeType.JAVASCRIPT:
            secure_env.update({
                'NODE_ENV': 'production',
                'NODE_OPTIONS': '--max-old-space-size=128 --max-semi-space-size=64',  # 限制Node.js堆内存
            })
        
        # Bash特定的安全环境变量
        elif request.code_type == CodeType.BASH:
            secure_env.update({
                'BASH': '/bin/bash',
                'SHELL': '/bin/bash',
            })
        
        # 只添加安全的用户环境变量
        for key, value in request.environment_vars.items():
            if self._is_safe_env_var(key, value):
                secure_env[key] = value
        
        return secure_env
    
    def _is_safe_env_var(self, key: str, value: str) -> bool:
        """检查环境变量是否安全"""
        # 危险的环境变量
        dangerous_vars = {
            'PATH', 'LD_LIBRARY_PATH', 'LD_PRELOAD', 'PYTHONPATH',
            'NODE_PATH', 'HOME', 'SHELL', 'USER', 'LOGNAME'
        }
        
        if key in dangerous_vars:
            return False
        
        # 检查值是否包含危险内容
        if any(char in value for char in ['`', '$', ';', '|', '&', '>', '<']):
            return False
        
        return True
    
    async def execute_code(self, request: CodeExecutionRequest, env_info: Dict[str, Any], security_config: Dict[str, Any]) -> ExecutionResult:
        """安全地执行代码"""
        code_file = env_info["code_file"]
        work_dir = env_info["work_dir"]
        
        # 获取执行命令
        runner_func = self.code_runners.get(request.code_type)
        if not runner_func:
            # 注意：task_id由基类的run_task方法设置，这里保持为空字符串
            return ExecutionResult(
                task_id="",
                status=TaskStatus.FAILED,
                error=f"不支持的代码类型: {request.code_type}",
                output=None,
                return_code=1,
                execution_time=0.0,
                memory_usage=0,
                cpu_usage=0.0,
                start_time=datetime.now(),
                end_time=datetime.now()
            )
        
        # 获取输入参数
        input_arg = json.dumps(request.input_arg) if request.input_arg else ""
        cmd = runner_func(code_file, request.timeout_seconds, input_arg)
        start_time = datetime.now()
        
        try:
            # 创建安全的环境
            
            secure_env = self._create_secure_environment(request, work_dir)
            
            # 使用更严格的执行方式
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=work_dir,
                env=secure_env,
                preexec_fn=self._create_secure_preexec_fn(security_config, request.code_type),
                # 额外的安全措施
                start_new_session=True  # 创建新的会话
            )
            
            # 等待执行完成，JavaScript需要更长的超时时间
            if request.code_type == CodeType.JAVASCRIPT:
                timeout = min(request.timeout_seconds, 30)  # JavaScript允许更长超时
            else:
                timeout = min(request.timeout_seconds, 15)  # 其他语言稍短超时
            try:
                stdout, stderr = await asyncio.wait_for(
                    process.communicate(),
                    timeout=timeout
                )
            except asyncio.TimeoutError:
                # 强制终止进程
                await self._kill_process_group(process)
                raise
            
            end_time = datetime.now()
            execution_time = (end_time - start_time).total_seconds()
            
            # 限制输出大小
            max_output_size = 4096  # 4KB
            output_text = stdout.decode('utf-8', errors='replace')[:max_output_size] if stdout is not None else None
            error_text = stderr.decode('utf-8', errors='replace')[:max_output_size] if stderr is not None else None
            
            # 简化的资源使用统计
            memory_usage = 0
            cpu_usage = 0.0
            
            return ExecutionResult(
                task_id="",
                status=TaskStatus.COMPLETED if process.returncode == 0 else TaskStatus.FAILED,
                output=output_text,
                error=error_text,
                return_code=process.returncode,
                memory_usage=memory_usage,
                cpu_usage=cpu_usage,
                execution_time=execution_time,
                start_time=start_time,
                end_time=end_time
            )
            
        except asyncio.TimeoutError:
            await self._kill_process_group(process)
            return ExecutionResult(
                task_id="",
                status=TaskStatus.TIMEOUT,
                error="代码执行超时",
                output=None,
                return_code=-1,
                execution_time=(datetime.now() - start_time).total_seconds(),
                memory_usage=0,
                cpu_usage=0.0,
                start_time=start_time,
                end_time=datetime.now()
            )
            
        except Exception as e:
            logger.error(f"代码执行失败: {e}")
            return ExecutionResult(
                task_id="",
                status=TaskStatus.FAILED,
                error=f"执行出错: {str(e)}",
                output=None,
                return_code=-1,
                execution_time=(datetime.now() - start_time).total_seconds(),
                memory_usage=0,
                cpu_usage=0.0,
                start_time=start_time,
                end_time=datetime.now()
            )
    
    async def _kill_process_group(self, process):
        """安全地终止进程组"""
        try:
            if process.pid:
                # 尝试终止整个进程组
                os.killpg(os.getpgid(process.pid), signal.SIGTERM)
                await asyncio.sleep(0.1)
                
                # 如果还在运行，强制杀死
                if process.returncode is None:
                    os.killpg(os.getpgid(process.pid), signal.SIGKILL)
        except (ProcessLookupError, PermissionError, OSError):
            pass
    
    async def cleanup_environment(self, env_info: Dict[str, Any]) -> None:
        """清理执行环境"""
        temp_dir = env_info.get("temp_dir")
        if temp_dir and os.path.exists(temp_dir):
            try:
                # 确保所有进程都被终止
                await asyncio.sleep(0.1)
                
                # 递归删除临时目录
                shutil.rmtree(temp_dir, ignore_errors=True)
                logger.debug(f"已清理临时目录: {temp_dir}")
            except Exception as e:
                logger.warning(f"清理临时目录失败: {e}")
        
        # 从字典中移除
        request_id = str(id(env_info))
        if request_id in self.temp_dirs:
            del self.temp_dirs[request_id] 